package com.hys.ai.controller;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.RequestResponseAdvisor;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;

/**
 * @author hys
 * @version Id: TongyiSimpleController.java, v 0.1 2024年07月31日 09:22 hys Exp $
 */
@Slf4j
@RestController
public class RagController {

    @Autowired
    private VectorStore vectorStore;

    @Autowired
    private       ChatModel  chatModel;
    // 模拟数据库存储会话和消息
    private final ChatMemory chatMemory = new InMemoryChatMemory();

    /**
     * 从向量数据库中查找文档，并将查询的文档作为上下文回答。
     *
     * @param prompt 用户的提问
     * @return SSE流响应
     */
    @RequestMapping(value = "chat/stream/history", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> chatStreamWithDatabase(@RequestParam(defaultValue = "我是谁，我做过哪些项目") String prompt) {
        log.info("=========<"+prompt);
        // 1. 定义提示词模板，question_answer_context会被替换成向量数据库中查询到的文档。
        String promptWithContext = """
                下面是上下文信息
                ---------------------
                {question_answer_context}
                ---------------------
                给定的上下文和提供的历史信息，而不是事先的知识，回复用户的意见。如果答案不在上下文中，告诉用户你不能回答这个问题。
                """;
        QuestionAnswerAdvisor questionAnswerAdvisor = new QuestionAnswerAdvisor(vectorStore, SearchRequest.defaults(), promptWithContext);
        //MessageChatMemoryAdvisor messageChatMemoryAdvisor = new MessageChatMemoryAdvisor(chatMemory);
        List<RequestResponseAdvisor> advisors = new ArrayList<>();
        advisors.add(questionAnswerAdvisor);
        //advisors.add(messageChatMemoryAdvisor);

        Flux<String> content = ChatClient.create(chatModel).prompt()
                .user(prompt)
                // 2. QuestionAnswerAdvisor会在运行时替换模板中的占位符`question_answer_context`，替换成向量数据库中查询到的文档。此时的query=用户的提问+替换完的提示词模板;
                .advisors(advisors)
                .stream()
                // 3. query发送给大模型得到答案
                .content();
        log.info("=========>"+content.toString());
        return content;


    }

}
